# 导入库
import pymysql
import pandas as pd
import pymysql.cursors


class MysqlUtils(object):
    """数据库连接工具类"""
    def __init__(self):
        # 初始化数据库连接
        self.conn = pymysql.connect(
            host='127.0.0.1',
            user='root',
            password='root',
            database='text',
            port=3306,
            charset='utf8'
        )


class ClassificationUtils(object):
    """数据分类处理工具类"""
    def __init__(self):
        pass

    def get_fina_indicator(self, conn):
        """
        获取财务数据
        :param conn: 数据库连接对象
        :return: 处理后的财务数据DataFrame
        """
        cursor = conn.cursor(cursor=pymysql.cursors.DictCursor)
        sql = """
        SELECT ts_code, ann_date, eps, total_revenue_ps, undist_profit_ps,
               gross_margin, fcff, fcfe, tangible_asset, bps, 
               grossprofit_margin, npta
        FROM financial_data 
        WHERE ann_date BETWEEN '2023-01-01' and '2024-01-01'
        """
        cursor.execute(sql)
        rows = cursor.fetchall()
        df = pd.DataFrame(rows)
        
        # 处理缺失数据（删除包含关键指标缺失的行）
        df = df.dropna(subset=[
            'eps', 'total_revenue_ps', 'undist_profit_ps', 'gross_margin', 
            'fcff', 'fcfe', 'tangible_asset', 'bps', 'grossprofit_margin', 'npta'
        ])
        
        # 重建索引（删除缺失行后重置索引）
        df = df.reset_index(drop=True)
        return df

    def get_daily(self, conn, df):
        """
        获取日线数据并与财务数据整合
        :param conn: 数据库连接对象
        :param df: 财务数据DataFrame
        """
        cursor = conn.cursor(cursor=pymysql.cursors.DictCursor)
        new_list = []  # 用于存储整合后的数据
        
        # 遍历财务数据每一行
        for index, row in df.iterrows():
            print(index)  # 打印当前处理的索引
            ts_code = row['ts_code']
            ann_date = row['ann_date'].strftime('%Y-%m-%d')  # 格式化公告日期
            
            # 查询公告日后的20条日线数据
            sql = f"""
            SELECT trade_date, closes 
            FROM date_1 d 
            WHERE d.the_date > Date('{ann_date}') 
              AND d.ts_code = '{ts_code}'
            ORDER BY d.trade_date ASC 
            LIMIT 20
            """
            cursor.execute(sql)
            rows = cursor.fetchall()
            df1 = pd.DataFrame(rows)  # 转为DataFrame处理
            
            try:
                # 若存在日线数据，计算关键价格指标并整合
                if len(df1) > 0:
                    max_closes = df1['closes'].max()        # 20天内最大收盘价
                    min_closes = df1['closes'].min()        # 20天内最小收盘价
                    the_closes = df1['closes'].iloc[-1]     # 20天内最新收盘价
                    
                    # 整合财务数据与价格指标
                    new_list.append({
                        'ts_code': ts_code,
                        'ann_date': ann_date,
                        'max_closes': max_closes,
                        'min_closes': min_closes,
                        'the_closes': the_closes,
                        'eps': row['eps'],
                        'total_revenue_ps': row['total_revenue_ps'],
                        'undist_profit_ps': row['undist_profit_ps'],
                        'gross_margin': row['gross_margin'],
                        'fcff': row['fcff'],
                        'fcfe': row['fcfe'],
                        'tangible_asset': row['tangible_asset'],
                        'bps': row['bps'],
                        'grossprofit_margin': row['grossprofit_margin'],
                        'npta': row['npta'],
                    })
            except Exception as e:
                print(f"处理异常: {e}")
                continue  # 捕获异常后继续处理下一行
        df2 = pd.DataFrame(new_list)
        df2.to_csv('fina_indicator.csv',index=False)


if __name__ == '__main__':
    # 主程序入口
    mu = MysqlUtils()  # 初始化数据库连接工具
    classification_utils = ClassificationUtils()  # 初始化数据处理工具
    
    # 执行数据处理流程
    df = classification_utils.get_fina_indicator(mu.conn)  # 获取财务数据
    classification_utils.get_daily(mu.conn, df)  # 处理日线数据并整合